import numpy as np

class GomokuEnv:
    """
    五子棋环境 - 增强版威胁检测和防守奖励系统
    State: 展平的棋盘状态，每个位置：1 (当前玩家), -1 (对手), 0 (空)
    Action: 0 到 board_size*board_size - 1 的整数，表示落子位置
    Reward: 细粒度奖励系统 + 增强威胁检测
    """

    def __init__(self, board_size=15):
        self.board_size = board_size
        self.observation_space_shape = (board_size * board_size,)
        self.action_space_n = board_size * board_size

        self.board = None
        self.current_player = None
        self.winner = None
        self.done = False

        # 威胁等级定义
        self.THREAT_LEVELS = {
            'CRITICAL': 10.0,    # 阻止对手立即获胜 (4子连线)
            'HIGH': 3.0,        # 阻止对手形成活四
            'MEDIUM': 1,      # 阻止对手形成活三
            'LOW': 0.5          # 阻止对手形成活二
        }

        self.reset()

    def reset(self):
        """重置棋盘状态，开始新游戏"""
        self.board = np.zeros((self.board_size, self.board_size), dtype=int)
        self.current_player = 1
        self.winner = None
        self.done = False
        return self._get_state()

    def step(self, action):
        """
        执行一个落子动作并更新环境状态。
        action: 0到board_size*board_size-1的整数，表示落子位置。
        """
        if self.done:
            return self._get_state(), 0, True, {"winner": self.winner, "game_over": True}

        row, col = divmod(action, self.board_size)

        # 检查非法落子
        if not (0 <= row < self.board_size and 0 <= col < self.board_size and self.board[row, col] == 0):
            reward = -20.0  # 非法落子惩罚
            self.done = True
            self.winner = -self.current_player
            return self._get_state(), reward, True, {"winner": self.winner, "invalid_move": True}

        # 执行落子
        self.board[row, col] = self.current_player

        # 细粒度奖励计算
        info = {"winner": None}

        # 检查是否获胜
        if self._check_win(row, col, self.current_player):
            reward = 20.0  # 获胜奖励
            self.winner = self.current_player
            self.done = True
            info["winner"] = self.winner
        # 检查是否平局
        elif self._is_board_full():
            reward = 0.0
            self.winner = 0
            self.done = True
            info["winner"] = 0
        else:
            # 🔥 使用增强的中间奖励计算
            reward = self._calculate_enhanced_intermediate_reward(row, col, self.current_player)
            self.current_player *= -1

            # 额外信息记录
            max_line = self._get_max_line_length(row, col, self.board[row, col])
            info["max_line"] = max_line

            # 🔥 新增：威胁检测信息
            threat_info = self._get_threat_analysis_info(row, col)
            info.update(threat_info)

        return self._get_state(), reward, self.done, info

    def _calculate_enhanced_intermediate_reward(self, r, c, player):
        """🔥 增强版中间奖励计算 - 整合威胁检测逻辑"""
        base_reward = -0.01  # 基础步骤惩罚

        # 自己的进攻奖励
        attack_reward = self._calculate_attack_reward(r, c, player)

        # 🔥 增强的防守奖励
        defense_reward = self._enhanced_threat_detection(r, c, player)

        # 位置价值奖励（中心区域更有价值）
        center = self.board_size // 2
        distance_from_center = abs(r - center) + abs(c - center)
        position_reward = max(0, (center - distance_from_center) * 0.01)

        return base_reward + attack_reward + defense_reward + position_reward

    def _enhanced_threat_detection(self, r, c, player):
        """🔥 增强的威胁检测和阻止策略"""
        opponent = -player
        total_defense_reward = 0

        directions = [(0, 1), (1, 0), (1, 1), (1, -1)]

        for dr, dc in directions:
            threat_level = self._analyze_threat_in_direction(r, c, dr, dc, opponent)
            total_defense_reward += self.THREAT_LEVELS.get(threat_level, 0)

        # 检测多重威胁情况
        multiple_threats = self._detect_multiple_threats(r, c, opponent)
        if multiple_threats > 1:
            total_defense_reward *= (1 + multiple_threats * 0.5)  # 多重威胁加成

        return total_defense_reward

    def _analyze_threat_in_direction(self, r, c, dr, dc, opponent):
        """🔥 分析特定方向的威胁等级"""
        # 扩大检测范围到7个位置
        positions = []
        coords = []

        for i in range(-3, 4):
            nr, nc = r + dr * i, c + dc * i
            if 0 <= nr < self.board_size and 0 <= nc < self.board_size:
                positions.append(self.board[nr, nc])
                coords.append((nr, nc))
            else:
                positions.append(-999)  # 边界标记
                coords.append(None)

        # 当前位置在中心 (index 3)
        positions[3] = opponent  # 假设对手不在这里下，我们下了会如何

        # 分析威胁模式
        threat_patterns = self._identify_threat_patterns(positions, opponent)

        return self._evaluate_threat_level(threat_patterns)

    def _identify_threat_patterns(self, positions, opponent):
        """🔥 识别威胁模式"""
        patterns = {
            'four_in_row': False,      # 四子连线 (立即获胜威胁)
            'open_four': False,        # 活四 (下一步必胜)
            'closed_four': False,      # 冲四 (一端被堵)
            'open_three': False,       # 活三 (可形成活四)
            'closed_three': False,     # 眠三 (一端被堵)
            'open_two': False,         # 活二
            'closed_two': False        # 眠二
        }

        # 检查连续棋子
        consecutive_counts = []
        current_count = 0

        for i, pos in enumerate(positions):
            if pos == opponent:
                current_count += 1
            else:
                if current_count > 0:
                    # 记录连续段的信息：(长度, 起始位置, 结束位置)
                    consecutive_counts.append({
                        'length': current_count,
                        'start': i - current_count,
                        'end': i - 1,
                        'left_empty': i - current_count - 1 >= 0 and positions[i - current_count - 1] == 0,
                        'right_empty': i < len(positions) and positions[i] == 0
                    })
                current_count = 0

        # 处理结尾的连续段
        if current_count > 0:
            consecutive_counts.append({
                'length': current_count,
                'start': len(positions) - current_count,
                'end': len(positions) - 1,
                'left_empty': len(positions) - current_count - 1 >= 0 and positions[len(positions) - current_count - 1] == 0,
                'right_empty': False
            })

        # 根据连续段判断威胁模式
        for segment in consecutive_counts:
            length = segment['length']
            left_empty = segment['left_empty']
            right_empty = segment['right_empty']

            if length >= 4:
                patterns['four_in_row'] = True
            elif length == 3:
                if left_empty and right_empty:
                    patterns['open_four'] = True  # 两端都空，可形成活四
                elif left_empty or right_empty:
                    patterns['closed_four'] = True  # 一端空，冲四
            elif length == 2:
                if left_empty and right_empty:
                    patterns['open_three'] = True
                elif left_empty or right_empty:
                    patterns['closed_three'] = True
            elif length == 1:
                if left_empty and right_empty:
                    patterns['open_two'] = True
                elif left_empty or right_empty:
                    patterns['closed_two'] = True

        return patterns

    def _evaluate_threat_level(self, patterns):
        """🔥 评估威胁等级"""
        if patterns['four_in_row']:
            return 'CRITICAL'
        elif patterns['open_four']:
            return 'HIGH'
        elif patterns['closed_four'] or patterns['open_three']:
            return 'MEDIUM'
        elif patterns['closed_three'] or patterns['open_two']:
            return 'LOW'
        else:
            return None

    def _detect_multiple_threats(self, r, c, opponent):
        """🔥 检测多重威胁情况"""
        threat_count = 0
        directions = [(0, 1), (1, 0), (1, 1), (1, -1)]

        for dr, dc in directions:
            threat_level = self._analyze_threat_in_direction(r, c, dr, dc, opponent)
            if threat_level in ['CRITICAL', 'HIGH', 'MEDIUM']:
                threat_count += 1

        return threat_count

    def _calculate_attack_reward(self, r, c, player):
        """🔥 计算进攻奖励"""
        max_line = self._get_max_line_length(r, c, player)

        attack_rewards = {
            5: 10.0,    # 获胜
            4: 3.0,     # 四子连线
            3: 1.0,     # 三子连线
            2: 0.2      # 二子连线
        }

        return attack_rewards.get(max_line, 0)

    def _get_threat_analysis_info(self, r, c):
        """🔥 获取威胁分析信息用于调试"""
        current_player = self.board[r, c]
        opponent = -current_player

        info = {
            "threats_detected": {},
            "defense_reward": 0,
            "attack_reward": 0,
            "multiple_threats": 0
        }

        # 分析各方向威胁
        directions = [(0, 1), (1, 0), (1, 1), (1, -1)]
        direction_names = ["horizontal", "vertical", "diagonal1", "diagonal2"]

        for i, (dr, dc) in enumerate(directions):
            threat_level = self._analyze_threat_in_direction(r, c, dr, dc, opponent)
            if threat_level:
                info["threats_detected"][direction_names[i]] = threat_level

        info["defense_reward"] = self._enhanced_threat_detection(r, c, current_player)
        info["attack_reward"] = self._calculate_attack_reward(r, c, current_player)
        info["multiple_threats"] = self._detect_multiple_threats(r, c, opponent)

        return info

    def _get_max_line_length(self, r, c, player):
        """获取某个位置在所有方向上的最大连线长度"""
        directions = [(0, 1), (1, 0), (1, 1), (1, -1)]
        max_length = 0

        for dr, dc in directions:
            length = self._count_line(r, c, dr, dc, player)
            max_length = max(max_length, length)

        return max_length

    def _count_line(self, r, c, dr, dc, player):
        """计算从(r,c)位置沿(dr,dc)方向的连续棋子数"""
        count = 1

        # 正方向
        for i in range(1, 5):
            nr, nc = r + dr * i, c + dc * i
            if 0 <= nr < self.board_size and 0 <= nc < self.board_size and self.board[nr, nc] == player:
                count += 1
            else:
                break

        # 反方向
        for i in range(1, 5):
            nr, nc = r - dr * i, c - dc * i
            if 0 <= nr < self.board_size and 0 <= nc < self.board_size and self.board[nr, nc] == player:
                count += 1
            else:
                break

        return count

    def _count_line_with_gap(self, r, c, dr, dc, player):
        """计算包含一个空位的潜在连线长度（用于检测威胁）"""
        positions = []

        # 收集该方向上的5个位置
        for i in range(-2, 3):
            nr, nc = r + dr * i, c + dc * i
            if 0 <= nr < self.board_size and 0 <= nc < self.board_size:
                positions.append(self.board[nr, nc])
            else:
                positions.append(-999)  # 边界外

        # 计算player棋子数
        player_count = positions.count(player)
        empty_count = positions.count(0) + positions.count(-999)  # 当前位置算作空

        if empty_count == 1:  # 只有一个空位
            return player_count
        return 0

    def _get_state(self):
        """返回展平的棋盘状态"""
        if self.current_player == 1:
            return self.board.flatten().astype(np.float32)
        else:
            return (-self.board).flatten().astype(np.float32)

    def _check_win(self, r, c, player):
        """检查给定位置是否导致player获胜"""
        directions = [(0, 1), (1, 0), (1, 1), (1, -1)]

        for dr, dc in directions:
            count = 1
            for i in range(1, 5):
                nr, nc = r + dr * i, c + dc * i
                if 0 <= nr < self.board_size and 0 <= nc < self.board_size and self.board[nr, nc] == player:
                    count += 1
                else:
                    break
            for i in range(1, 5):
                nr, nc = r - dr * i, c - dc * i
                if 0 <= nr < self.board_size and 0 <= nc < self.board_size and self.board[nr, nc] == player:
                    count += 1
                else:
                    break
            if count >= 5:
                return True
        return False

    def _is_board_full(self):
        """检查棋盘是否已满"""
        return np.all(self.board != 0)

    def get_legal_actions(self):
        """返回所有合法动作的索引列表"""
        legal_actions = []
        for i in range(self.board_size):
            for j in range(self.board_size):
                if self.board[i, j] == 0:
                    legal_actions.append(i * self.board_size + j)
        return legal_actions
    def close(self):
        """关闭环境"""
        pass